package v1

import (
	"bytes"
	"context"
	"encoding/json"
	"html/template"
	"log/slog"
	"strings"
	"time"

	"quipus/internal/database/model"
	llm "quipus/internal/llm"
	"quipus/internal/server/process"

	"quipus/internal/llm/factory"
	PROMPT "quipus/internal/llm/prompt"
	"quipus/internal/utils"
	v1pb "quipus/proto/gen/api/v1"
	apb "quipus/proto/gen/assist"
	cpb "quipus/proto/gen/common"

	"quipus/internal/llm/provider"

	llmMaker "quipus/internal/llm/factory"

	"github.com/gofiber/websocket/v2"
	"github.com/google/cel-go/cel"
	"github.com/lithammer/shortuuid/v4"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
	expr "google.golang.org/genproto/googleapis/api/expr/v1alpha1"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/emptypb"
	"gorm.io/gorm"
)

func (s *APIV1Service) CreateLLMProvider(ctx context.Context, request *v1pb.CreateLLMProviderRequest) (*v1pb.LLMProvider, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	provider := convertLLMProviderToDB(request)
	provider.TenantID = currentUser.TenantID
	provider.CreatorID = currentUser.ID
	err = s.Gds.Handler.CreateLLMProvider(ctx, provider)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to create llm provider, error: %+v", err)
	}
	return convertLLMProviderFromDB(provider), nil
}

func convertLLMProviderFromDB(provider *model.LlmProvider) *v1pb.LLMProvider {
	temp := &v1pb.LLMProvider{
		Name:        provider.Name,
		ApiName:     provider.APIName,
		Description: provider.Description,
		AccessUrl:   provider.AccessURL,
		Urn:         utils.BuildUrn(utils.ProviderUrnPrefix, provider.ID),
		Tenant:      utils.BuildUrn(utils.TenantUrnPrefix, provider.TenantID),
		Provider:    provider.Provider,
	}

	if provider.AuthInfo != "" {
		key := []byte(viper.GetString("encryption-key"))
		auth, e := utils.Decrypt(provider.AuthInfo, key)
		if e != nil {
			return nil
		}
		var authInfo map[string]string
		e = json.Unmarshal([]byte(auth), &authInfo)
		if e != nil {
			return nil
		}

		for key, _ := range authInfo {
			authInfo[key] = "******"
		}
		temp.AuthInfo = authInfo
	}
	return temp
}

func convertLLMProviderToDB(req *v1pb.CreateLLMProviderRequest) *model.LlmProvider {
	temp := &model.LlmProvider{
		Name:        req.Name,
		APIName:     req.ApiName,
		AccessURL:   req.AccessUrl,
		Description: req.Description,
		Provider:    req.Provider,
	}

	if req.AuthInfo != nil || len(req.AuthInfo) > 1 {
		cb, e := json.Marshal(req.AuthInfo)
		if e != nil {
			return nil
		}
		temp.AuthInfo = string(cb)
		key := []byte(viper.GetString("encryption-key"))
		auth, e := utils.Encrypt(temp.AuthInfo, key)
		if e != nil {
			return nil
		}
		temp.AuthInfo = auth
	}
	return temp
}

func (s *APIV1Service) UpdateLLMProvider(ctx context.Context, request *v1pb.UpdateLLMProviderRequest) (*v1pb.LLMProvider, error) {
	providerId, err := utils.ExtractIDFromUrn(utils.ProviderUrnPrefix, request.Provider.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid user name: %v", err)
	}
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	if request.UpdateMask == nil || len(request.UpdateMask.Paths) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "update mask is empty")
	}

	provider, err := s.Gds.Handler.GetLLMProvider(ctx, providerId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if provider == nil {
		return nil, status.Errorf(codes.NotFound, "user not found")
	}

	if currentUser.Role == cpb.Role_TENANT_ADMIN.String() && provider.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	currentTs := time.Now().Unix()
	update := &apb.UpdateLLMProvider{
		Id:        provider.ID,
		UpdatedTs: currentTs,
	}

	for _, field := range request.UpdateMask.Paths {
		if field == "name" {
			update.Name = request.Provider.Name
		} else if field == "api" {
			update.Api = request.Provider.ApiName
		} else if field == "access_url" {
			update.Url = request.Provider.AccessUrl
		} else if field == "description" {
			update.Description = request.Provider.Description
		} else if field == "auth_info" {
			if request.Provider.AuthInfo != nil || len(request.Provider.AuthInfo) > 1 {
				cb, e := json.Marshal(request.Provider.AuthInfo)
				if e != nil {
					return nil, nil
				}
				a := string(cb)
				key := []byte(viper.GetString("encryption-key"))
				auth, e := utils.Encrypt(a, key)
				if e != nil {
					return nil, nil
				}
				update.AuthInfo = auth
			}
		} else {
			return nil, status.Errorf(codes.InvalidArgument, "invalid update path: %s", field)
		}
	}

	_, err = s.Gds.Handler.UpdateLLMProvider(ctx, update)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to update user: %v", err)
	}

	updated, err := s.Gds.Handler.GetLLMProvider(ctx, providerId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	return convertLLMProviderFromDB(updated), nil
}

func (s *APIV1Service) DeleteLLMProvider(ctx context.Context, request *v1pb.DeleteLLMProviderRequest) (*emptypb.Empty, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	id, err := utils.ExtractIDFromUrn(utils.ProviderUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	provider, err := s.Gds.Handler.GetLLMProvider(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get llm provider, error: %+v", err)
	}

	if provider == nil {
		return nil, status.Errorf(codes.NotFound, "llm provider not found")
	}

	if currentUser.Role == cpb.Role_TENANT_ADMIN.String() && provider.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	err = s.Gds.Handler.DeleteLLMProvider(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete llm provider, error: %+v", err)
	}
	return nil, nil
}

type ProviderFilter struct {
	api  string
	name string
}

func findProviderField(callExpr *expr.Expr_Call, filter *ProviderFilter) {
	if len(callExpr.Args) == 2 {
		idExpr := callExpr.Args[0].GetIdentExpr()
		if idExpr != nil {
			if idExpr.Name == "api" {
				filter.api = callExpr.Args[1].GetConstExpr().GetStringValue()
			} else if idExpr.Name == "name" {
				filter.name = callExpr.Args[1].GetConstExpr().GetStringValue()
			}
			return
		}
	}
	for _, arg := range callExpr.Args {
		callExpr := arg.GetCallExpr()
		if callExpr != nil {
			findProviderField(callExpr, filter)
		}
	}
}

var ProviderFilterCELAttributes = []cel.EnvOption{
	cel.Variable("api", cel.StringType),
	cel.Variable("name", cel.StringType),
}

func parseProviderFilter(expression string) (*ProviderFilter, error) {
	e, err := cel.NewEnv(ProviderFilterCELAttributes...)
	if err != nil {
		return nil, err
	}
	ast, issues := e.Compile(expression)
	if issues != nil {
		return nil, errors.Errorf("found issue %v", issues)
	}
	filter := &ProviderFilter{}
	expr, err := cel.AstToParsedExpr(ast)
	if err != nil {
		return nil, err
	}
	callExpr := expr.GetExpr().GetCallExpr()
	findProviderField(callExpr, filter)
	return filter, nil
}

func buildProviderFindWithFilter(find *apb.FindLLMProvider, filter string) error {
	if find == nil {
		find = &apb.FindLLMProvider{}
	}

	if filter != "" {
		pfilter, err := parseProviderFilter(filter)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid filter: %v", err)
		}
		if len(pfilter.api) > 0 {
			find.Api = pfilter.api
		}

		if len(pfilter.name) > 0 {
			find.Name = pfilter.name
		}
	}

	return nil
}

func (s *APIV1Service) ListLLMProviders(ctx context.Context, request *v1pb.ListLLMProviderRequest) (*v1pb.ListLLMProviderResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	find := &apb.FindLLMProvider{}
	buildProviderFindWithFilter(find, request.Filter)

	var limit, offset int32
	if request.PageToken != nil {
		limit = request.PageToken.Limit
		offset = request.PageToken.Offset
	} else {
		limit = request.PageSize
	}
	if limit <= 0 {
		limit = DefaultPageSize
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		find.TenantId = currentUser.TenantID
	}

	limitPlusOne := limit
	find.Limit = limitPlusOne
	find.Offset = offset

	providers, count, err := s.Gds.Handler.ListLLMProviders(ctx, find)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list subject types: %v", err)
	}
	providerMessages := []*v1pb.LLMProvider{}

	if int32(len(providers)) == limitPlusOne {
		providers = providers[:limit]
	}

	for _, p := range providers {
		pMsg := convertLLMProviderFromDB(p)
		providerMessages = append(providerMessages, pMsg)
	}

	response := &v1pb.ListLLMProviderResponse{
		TotalSize: count,
		Providers: providerMessages,
		NextPageToken: &cpb.PageToken{
			Limit:  int32(limit),
			Offset: int32(offset + limit),
		},
	}

	return response, nil
}

func (s *APIV1Service) CreateLLMModel(ctx context.Context, request *v1pb.CreateLLMModelRequest) (*v1pb.LLMModel, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	if request.Sp == "" {
		return nil, status.Errorf(codes.InvalidArgument, "service provider is required")
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	model := convertLLMModelToDB(request)
	model.CreatorID = currentUser.ID
	model.TenantID = currentUser.TenantID
	spId, err := utils.ExtractIDFromUrn(utils.ProviderUrnPrefix, request.Sp)
	model.ProviderID = spId
	err = s.Gds.Handler.CreateLLMModel(ctx, model)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to create llm model, error: %+v", err)
	}
	return convertLLMModelFromDB(model), nil
}

func convertLLMModelToDB(req *v1pb.CreateLLMModelRequest) *model.LlmModel {
	temp := &model.LlmModel{
		Name:        req.Name,
		Model:       req.Model,
		ModelType:   req.ModelType.String(),
		Description: req.Description,
	}

	if req.Config != nil || len(req.Config) > 1 {
		cb, e := json.Marshal(req.Config)
		if e != nil {
			return nil
		}
		temp.ConfigInfo = string(cb)
	}

	if req.Options != nil || len(req.Options) > 1 {
		cb, e := json.Marshal(req.Options)
		if e != nil {
			return nil
		}
		temp.OptionInfo = string(cb)
	}
	return temp
}

func convertLLMModelFromDB(m *model.LlmModel) *v1pb.LLMModel {
	temp := &v1pb.LLMModel{
		Name:        m.Name,
		Model:       m.Model,
		ModelType:   v1pb.ModelType(v1pb.ModelType_value[m.ModelType]),
		Description: m.Description,
		Urn:         utils.BuildUrn(utils.ModelUrnPrefix, m.ID),
		Sp:          utils.BuildUrn(utils.ProviderUrnPrefix, m.ProviderID),
		Tenant:      utils.BuildUrn(utils.TenantUrnPrefix, m.TenantID),
	}

	if m.ConfigInfo != "" {

		var c map[string]string
		e := json.Unmarshal([]byte(m.ConfigInfo), &c)
		if e != nil {
			return nil
		}
		temp.Config = c
	}

	if m.OptionInfo != "" {

		var o map[string]string
		e := json.Unmarshal([]byte(m.OptionInfo), &o)
		if e != nil {
			return nil
		}
		temp.Options = o
	}
	return temp
}

func (s *APIV1Service) UpdateLLMModel(ctx context.Context, request *v1pb.UpdateLLMModelRequest) (*v1pb.LLMModel, error) {
	modelId, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, request.Model.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid user name: %v", err)
	}
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	if request.UpdateMask == nil || len(request.UpdateMask.Paths) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "update mask is empty")
	}

	model, err := s.Gds.Handler.GetLLMModel(ctx, modelId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get model: %v", err)
	}
	if model == nil {
		return nil, status.Errorf(codes.NotFound, "model not found")
	}

	if currentUser.Role == cpb.Role_TENANT_ADMIN.String() && model.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	currentTs := time.Now().Unix()
	update := &apb.UpdateLLMModel{
		Id:        model.ID,
		UpdatedTs: currentTs,
	}

	for _, field := range request.UpdateMask.Paths {
		if field == "name" {
			update.Name = request.Model.Name
		} else if field == "model" {
			update.Model = request.Model.Model
		} else if field == "description" {
			update.Description = request.Model.Description
		} else if field == "config" {
			if request.Model.Config != nil || len(request.Model.Config) > 1 {
				cb, e := json.Marshal(request.Model.Config)
				if e != nil {
					return nil, nil
				}
				c := string(cb)
				update.Config = c
			}
		} else if field == "options" {
			if request.Model.Options != nil || len(request.Model.Options) > 1 {
				cb, e := json.Marshal(request.Model.Options)
				if e != nil {
					return nil, nil
				}
				o := string(cb)
				update.Options = o
			}
		} else {
			return nil, status.Errorf(codes.InvalidArgument, "invalid update path: %s", field)
		}
	}

	_, err = s.Gds.Handler.UpdateLLMModel(ctx, update)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to update user: %v", err)
	}

	updated, err := s.Gds.Handler.GetLLMModel(ctx, modelId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	return convertLLMModelFromDB(updated), nil
}

func (s *APIV1Service) DeleteLLMModel(ctx context.Context, request *v1pb.DeleteLLMModelRequest) (*emptypb.Empty, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	id, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	m, err := s.Gds.Handler.GetLLMModel(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get llm model, error: %+v", err)
	}

	if m == nil {
		return nil, status.Errorf(codes.NotFound, "model not found")
	}

	if currentUser.Role == cpb.Role_TENANT_ADMIN.String() && m.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	err = s.Gds.Handler.DeleteLLMModel(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete llm model, error: %+v", err)
	}
	return nil, nil
}

type ModelFilter struct {
	model      string
	name       string
	providerId int32
	modelType  string
}

func findModelField(callExpr *expr.Expr_Call, filter *ModelFilter) {
	if len(callExpr.Args) == 2 {
		idExpr := callExpr.Args[0].GetIdentExpr()
		if idExpr != nil {
			if idExpr.Name == "model" {
				filter.model = callExpr.Args[1].GetConstExpr().GetStringValue()
			} else if idExpr.Name == "name" {
				filter.name = callExpr.Args[1].GetConstExpr().GetStringValue()
			} else if idExpr.Name == "model_type" {
				filter.modelType = callExpr.Args[1].GetConstExpr().GetStringValue()
			} else if idExpr.Name == "provider" {
				urn := callExpr.Args[1].GetConstExpr().GetStringValue()
				id, e := utils.ExtractIDFromUrn(utils.ProviderUrnPrefix, urn)
				if e != nil {
					return
				}

				filter.providerId = id
			}
			return
		}
	}

	for _, arg := range callExpr.Args {
		callExpr := arg.GetCallExpr()
		if callExpr != nil {
			findModelField(callExpr, filter)
		}
	}
}

var ModelFilterCELAttributes = []cel.EnvOption{
	cel.Variable("model", cel.StringType),
	cel.Variable("name", cel.StringType),
	cel.Variable("provider", cel.StringType),
	cel.Variable("model_type", cel.StringType),
}

func parseModelFilter(expression string) (*ModelFilter, error) {
	e, err := cel.NewEnv(ModelFilterCELAttributes...)
	if err != nil {
		return nil, err
	}
	ast, issues := e.Compile(expression)
	if issues != nil {
		return nil, errors.Errorf("found issue %v", issues)
	}
	filter := &ModelFilter{}
	expr, err := cel.AstToParsedExpr(ast)
	if err != nil {
		return nil, err
	}
	callExpr := expr.GetExpr().GetCallExpr()
	findModelField(callExpr, filter)
	return filter, nil
}

func buildModelFindWithFilter(find *apb.FindLLMModel, filter string) error {
	if find == nil {
		find = &apb.FindLLMModel{}
	}

	if filter != "" {
		mfilter, err := parseModelFilter(filter)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid filter: %v", err)
		}
		if len(mfilter.model) > 0 {
			find.Model = mfilter.model
		}

		if len(mfilter.name) > 0 {
			find.Name = mfilter.name
		}

		if len(mfilter.modelType) > 0 {
			find.ModelType = strings.ToUpper(mfilter.modelType)
		}

		if mfilter.providerId > 0 {
			find.ProviderId = mfilter.providerId
		}
	}

	return nil
}

func (s *APIV1Service) ListLLMModels(ctx context.Context, request *v1pb.ListLLMModelRequest) (*v1pb.ListLLMModelResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	find := &apb.FindLLMModel{}
	buildModelFindWithFilter(find, request.Filter)

	var limit, offset int32
	if request.PageToken != nil {
		limit = request.PageToken.Limit
		offset = request.PageToken.Offset
	} else {
		limit = request.PageSize
	}
	if limit <= 0 {
		limit = DefaultPageSize
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		find.TenantId = currentUser.TenantID
	}

	limitPlusOne := limit
	find.Limit = limitPlusOne
	find.Offset = offset

	models, count, err := s.Gds.Handler.ListLLMModels(ctx, find)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list subject types: %v", err)
	}
	modelMessages := []*v1pb.LLMModel{}

	if int32(len(models)) == limitPlusOne {
		models = models[:limit]
	}

	for _, m := range models {
		pMsg := convertLLMModelFromDB(m)
		modelMessages = append(modelMessages, pMsg)
	}

	response := &v1pb.ListLLMModelResponse{
		TotalSize: count,
		Models:    modelMessages,
		NextPageToken: &cpb.PageToken{
			Limit:  int32(limit),
			Offset: int32(offset + limit),
		},
	}

	return response, nil
}

func (s *APIV1Service) GeneratePrompt(ctx context.Context, request *v1pb.GeneratePromptRequest) (*v1pb.GeneratePromptResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	modelId, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, request.Model)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid model urn: %v", err)
	}

	model, err := s.Gds.Handler.GetLLMModel(ctx, modelId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get model: %v", err)
	}

	if model.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	p, err := s.Gds.Handler.GetLLMProvider(ctx, model.ProviderID)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get provider: %v", err)
	}

	config := &factory.ChatClientConfig{
		API:   p.APIName,
		Url:   p.AccessURL,
		Model: model.Model,
	}
	if p.AuthInfo != "" && len(p.AuthInfo) > 0 {
		key := []byte(viper.GetString("encryption-key"))
		authInfo, e := utils.Decrypt(p.AuthInfo, key)
		if e != nil {
			return nil, status.Errorf(codes.Internal, "failed to decrypt auth info: %v", e)
		}
		config.AuthInfo = authInfo
	}

	prompt, e := buildPrompt(request.Instruction)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to build prompt: %v", e)
	}

	llmClient := factory.NewChatClient(config)
	var options map[string]interface{}
	if request.Options != nil {
		e := json.Unmarshal([]byte(model.OptionInfo), &options)
		if e != nil {
			return nil, status.Errorf(codes.Internal, "failed to unmarshal options: %v", e)
		}

	}
	// TODO: call provider's API to generate prompt
	llmPrompt, err := queryPrompt(ctx, llmClient, prompt, options)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to query prompt: %v", err)
	}

	response := &v1pb.GeneratePromptResponse{
		Prompt: llmPrompt,
	}

	return response, nil
}

func buildPrompt(content string) (string, error) {
	type Prompt struct {
		Task_description string
	}

	Prompts := Prompt{
		Task_description: content,
	}

	tmpl, err := template.New("prompt").Parse(PROMPT.RULE_RAG_GENERATE_TEMPLATE)
	if err != nil {
		return "", err
	}

	buf := bytes.Buffer{}
	err = tmpl.Execute(&buf, Prompts)
	if err != nil {
		return "", err
	}

	return buf.String(), nil
}

func queryPrompt(ctx context.Context, client llm.LLM_Chat, prompt string, options map[string]interface{}) (string, error) {
	if len(prompt) < 1 {
		return "", errors.New("prompt is empty")
	}
	msgs := []llm.RoleContent{}

	msgs = append(msgs, llm.RoleContent{Role: llm.RoleUser, Content: prompt})

	req := &llm.CompleteReq{Msgs: msgs, Options: options}
	rsp, err := client.Complete(ctx, req)
	if err != nil {
		return "", err
	}
	if rsp == nil {
		return "", errors.New("response is nil")
	}
	return rsp.GetContent(), nil
}

func convertLLMAssistantToDB(req *v1pb.CreateLLMAssistantRequest) *model.LlmAssistant {
	id, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, req.Model)
	if err != nil {
		slog.Error("failed to get model id")
		return nil
	}

	temp := &model.LlmAssistant{
		Name:         req.Name,
		SystemPrompt: req.Prompt,
		ModelID:      id,
		Description:  req.Description,
		AvatarURL:    req.Avatar,
	}
	return temp
}

func convertLLMAssistantFromDB(assistant *model.LlmAssistant) *v1pb.LLMAssistant {
	temp := &v1pb.LLMAssistant{
		Name:        assistant.Name,
		Model:       utils.BuildUrn(utils.ModelUrnPrefix, assistant.ModelID),
		Description: assistant.Description,
		Prompt:      assistant.SystemPrompt,
		Avatar:      assistant.AvatarURL,
		Urn:         utils.BuildUrn(utils.AssistantUrnPrefix, assistant.ID),
		Tenant:      utils.BuildUrn(utils.TenantUrnPrefix, assistant.TenantID),
	}
	return temp
}

func (s *APIV1Service) CreateLLMAssistant(ctx context.Context, request *v1pb.CreateLLMAssistantRequest) (*v1pb.LLMAssistant, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	assistant := convertLLMAssistantToDB(request)
	assistant.TenantID = currentUser.TenantID
	assistant.CreatorID = currentUser.ID

	err = s.Gds.Handler.CreateLLMAssistant(ctx, assistant)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to create llm assistant, error: %+v", err)
	}
	assistantMsg := convertLLMAssistantFromDB(assistant)
	if request.VirtUser == true {

		user := &model.User{
			Username:     utils.HashString(assistantMsg.Urn),
			Role:         cpb.Role_AI_ASSISTANT.String(),
			Nickname:     assistant.Name,
			AvatarURL:    assistant.AvatarURL,
			PasswordHash: "xxxxxx",
			CreatorID:    currentUser.ID,
			TenantID:     currentUser.TenantID,
		}

		err = s.Gds.Handler.CreateUser(ctx, user)
		if err != nil {
			if err.Error() == gorm.ErrDuplicatedKey.Error() {
				slog.Info("user already exists", "user", user.Username)
				user, err = s.Gds.Handler.GetUser(ctx, &apb.FindUser{Username: assistant.Name,
					Role:      cpb.Role_AI_ASSISTANT,
					TenantId:  currentUser.TenantID,
					CreatorId: currentUser.ID})
				if err != nil {
					return nil, status.Errorf(codes.Internal, "ai user already exists, error: %+v", err)
				}
			} else {
				return nil, status.Errorf(codes.Internal, "failed to create ai user, error: %+v", err)
			}
		}

		virt := &model.VirtUser{
			ID:       user.ID,
			VirtID:   assistant.ID,
			VirtType: v1pb.ArtifactType_ARTIFACT_ASSISTANT.String(),
		}

		err = s.Gds.Handler.UpsertVirtUser(ctx, virt)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to create ai user, error: %+v", err)
		}
	}
	return assistantMsg, nil
}

func (s *APIV1Service) UpdateLLMAssistant(ctx context.Context, request *v1pb.UpdateLLMAssistantRequest) (*v1pb.LLMAssistant, error) {
	assistantId, err := utils.ExtractIDFromUrn(utils.AssistantUrnPrefix, request.Assistant.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid user name: %v", err)
	}
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	if request.UpdateMask == nil || len(request.UpdateMask.Paths) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "update mask is empty")
	}

	assistant, err := s.Gds.Handler.GetLLMAssistant(ctx, assistantId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if assistant == nil {
		return nil, status.Errorf(codes.NotFound, "user not found")
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		if assistant.TenantID != currentUser.TenantID {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}

		if assistant.CreatorID == currentUser.ID || currentUser.Role == cpb.Role_TENANT_ADMIN.String() {
			slog.Info("update assistant", "name", assistant.Name)
		} else {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}
	}

	if currentUser.Role == cpb.Role_TENANT_ADMIN.String() && assistant.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	currentTs := time.Now().Unix()
	update := &apb.UpdateLLMAssistant{
		Id:        assistantId,
		UpdatedTs: currentTs,
	}

	updateName := false
	for _, field := range request.UpdateMask.Paths {
		if field == "name" {
			update.Name = request.Assistant.Name
			updateName = true
		} else if field == "model" {
			id, e := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, request.Assistant.Model)
			if e != nil {
				return nil, status.Errorf(codes.InvalidArgument, "invalid model id: %s", request.Assistant.Model)
			}
			update.ModelId = id
		} else if field == "prompt" {
			if request.Assistant.Prompt == "" {
				return nil, status.Errorf(codes.InvalidArgument, "prompt is empty")
			}
			update.Prompt = request.Assistant.Prompt
		} else if field == "description" {
			update.Description = request.Assistant.Description
		} else if field == "virt_user" {
			user := &model.User{
				Username:     utils.HashString(request.Assistant.Urn),
				Nickname:     assistant.Name,
				Role:         cpb.Role_AI_ASSISTANT.String(),
				PasswordHash: "xxxxxx",
				CreatorID:    currentUser.ID,
				TenantID:     currentUser.TenantID,
			}

			err = s.Gds.Handler.CreateUser(ctx, user)
			if err != nil {
				if err.Error() == gorm.ErrDuplicatedKey.Error() {
					slog.Info("user already exists", "user", user.Username)
					user, err = s.Gds.Handler.GetUser(ctx, &apb.FindUser{Username: user.Username,
						Role:      cpb.Role_AI_ASSISTANT,
						TenantId:  currentUser.TenantID,
						CreatorId: currentUser.ID})
					if err != nil {
						return nil, status.Errorf(codes.Internal, "ai user already exists, error: %+v", err)
					}

					if updateName {
						s.Gds.Handler.UpdateUser(ctx, &apb.UpdateUser{
							Id:       user.ID,
							Nickname: assistant.Name,
						})
					}
				} else {
					return nil, status.Errorf(codes.Internal, "failed to create ai user, error: %+v", err)
				}
			}

			virt := &model.VirtUser{
				ID:       user.ID,
				VirtID:   assistant.ID,
				VirtType: v1pb.ArtifactType_ARTIFACT_ASSISTANT.String(),
			}

			err = s.Gds.Handler.UpsertVirtUser(ctx, virt)
			if err != nil {
				return nil, status.Errorf(codes.Internal, "failed to create ai user, error: %+v", err)
			}
		} else {
			return nil, status.Errorf(codes.InvalidArgument, "invalid update path: %s", field)
		}
	}

	_, err = s.Gds.Handler.UpdateLLMAssistant(ctx, update)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to update user: %v", err)
	}

	updated, err := s.Gds.Handler.GetLLMAssistant(ctx, assistantId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	return convertLLMAssistantFromDB(updated), nil
}

func (s *APIV1Service) DeleteLLMAssistant(ctx context.Context, request *v1pb.DeleteLLMAssistantRequest) (*emptypb.Empty, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() && currentUser.Role != cpb.Role_TENANT_ADMIN.String() {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	id, err := utils.ExtractIDFromUrn(utils.AssistantUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	assistant, err := s.Gds.Handler.GetLLMAssistant(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get llm assistant, error: %+v", err)
	}

	if assistant == nil {
		return nil, status.Errorf(codes.NotFound, "llm assistant not found")
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		if assistant.TenantID != currentUser.TenantID {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}

		if assistant.CreatorID == currentUser.ID || currentUser.Role == cpb.Role_TENANT_ADMIN.String() {
			slog.Info("delete assistant", "name", assistant.Name)
		} else {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}
	}

	err = s.Gds.Handler.DeleteLLMAssistant(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete llm provider, error: %+v", err)
	}

	delUser := &apb.DeleteUser{
		Id: id,
	}

	err = s.Gds.Handler.DeleteUser(ctx, delUser)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete user, error: %+v", err)
	}
	return nil, nil
}

func (s *APIV1Service) GetLLMAssistant(ctx context.Context, request *v1pb.GetLLMAssistantRequest) (*v1pb.LLMAssistant, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	id, err := utils.ExtractIDFromUrn(utils.AssistantUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	assistant, err := s.Gds.Handler.GetLLMAssistant(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get llm assistant, error: %+v", err)
	}

	if assistant == nil {
		return nil, status.Errorf(codes.NotFound, "llm assistant not found")
	}

	if currentUser.TenantID > 0 && assistant.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	aMsg := convertLLMAssistantFromDB(assistant)
	u, _ := s.Gds.Handler.GetUser(ctx, &apb.FindUser{Id: assistant.CreatorID})
	if u == nil {
		return nil, status.Errorf(codes.Internal, "failed to get user")
	}
	aMsg.Creator = s.convertUserFromStore(u)
	return aMsg, nil
}

type AssistantFilter struct {
	modelId int32
	name    string
}

var AssistantFilterCELAttributes = []cel.EnvOption{
	cel.Variable("model", cel.StringType),
	cel.Variable("name", cel.StringType),
}

func findAssistantField(callExpr *expr.Expr_Call, filter *AssistantFilter) {
	if len(callExpr.Args) == 2 {
		idExpr := callExpr.Args[0].GetIdentExpr()
		if idExpr != nil {
			if idExpr.Name == "name" {
				filter.name = callExpr.Args[1].GetConstExpr().GetStringValue()
			} else if idExpr.Name == "model" {
				m := callExpr.Args[1].GetConstExpr().GetStringValue()
				id, e := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, m)
				if e == nil {
					return
				}
				filter.modelId = id
			}
			return
		}
	}
	for _, arg := range callExpr.Args {
		callExpr := arg.GetCallExpr()
		if callExpr != nil {
			findAssistantField(callExpr, filter)
		}
	}
}

func parseAssistantFilter(expression string) (*AssistantFilter, error) {
	e, err := cel.NewEnv(AssistantFilterCELAttributes...)
	if err != nil {
		return nil, err
	}
	ast, issues := e.Compile(expression)
	if issues != nil {
		return nil, errors.Errorf("found issue %v", issues)
	}
	filter := &AssistantFilter{}
	expr, err := cel.AstToParsedExpr(ast)
	if err != nil {
		return nil, err
	}
	callExpr := expr.GetExpr().GetCallExpr()
	findAssistantField(callExpr, filter)
	return filter, nil
}

func buildAssistantFindWithFilter(find *apb.FindLLMAssistant, filter string) error {
	if find == nil {
		find = &apb.FindLLMAssistant{}
	}

	if filter != "" {
		afilter, err := parseAssistantFilter(filter)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid filter: %v", err)
		}
		if len(afilter.name) > 0 {
			find.Name = afilter.name
		}

		if afilter.modelId > 0 {
			find.ModelId = afilter.modelId
		}
	}

	return nil
}

func (s *APIV1Service) ListLLMAssistants(ctx context.Context, request *v1pb.ListLLMAssistantRequest) (*v1pb.ListLLMAssistantResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	find := &apb.FindLLMAssistant{}
	buildAssistantFindWithFilter(find, request.Filter)

	var limit, offset int32
	if request.PageToken != nil {
		limit = request.PageToken.Limit
		offset = request.PageToken.Offset
	} else {
		limit = request.PageSize
	}
	if limit <= 0 {
		limit = DefaultPageSize
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		find.TenantId = currentUser.TenantID
	}

	limitPlusOne := limit
	find.Limit = limitPlusOne
	find.Offset = offset

	assistants, count, err := s.Gds.Handler.ListLLMAssistants(ctx, find)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list assistants: %v", err)
	}
	assistantMessages := []*v1pb.LLMAssistant{}

	if int32(len(assistants)) == limitPlusOne {
		assistants = assistants[:limit]
	}

	userMap := make(map[int32]*v1pb.User)
	for _, a := range assistants {
		aMsg := convertLLMAssistantFromDB(a)
		user := userMap[a.CreatorID]
		if user == nil {
			u, _ := s.Gds.Handler.GetUser(ctx, &apb.FindUser{Id: a.CreatorID})
			if u == nil {
				continue
			}
			user = s.convertUserFromStore(u)

			if err != nil {
				return nil, errors.Wrap(err, "failed to convert user")
			}
		}

		aMsg.Creator = user

		userMap[a.CreatorID] = user
		assistantMessages = append(assistantMessages, aMsg)
	}

	response := &v1pb.ListLLMAssistantResponse{
		TotalSize:  count,
		Assistants: assistantMessages,
		NextPageToken: &cpb.PageToken{
			Limit:  int32(limit),
			Offset: int32(offset + limit),
		},
	}

	return response, nil
}

func (s *APIV1Service) ListLLMProviderSpecs(ctx context.Context, request *v1pb.ListLLMProviderSpecRequest) (*v1pb.ListLLMProviderSpecResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	locale := "cn"
	setting, err := s.Gds.Handler.GetUserSetting(ctx, currentUser.ID, apb.UserSettingKey_LOCALE.String())
	if err == nil {
		locale = setting.Value
	}

	providers, err := provider.ListProviderSpecs(request.Provider, request.Model, locale)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list provider specs: %v", err)
	}
	if providers == nil {
		return nil, status.Errorf(codes.Internal, "failed to list provider specs")
	}

	resp := &v1pb.ListLLMProviderSpecResponse{
		Providers: providers,
	}
	return resp, nil
}

func (s *APIV1Service) AssistantTest(ctx context.Context, request *v1pb.AssistantTestRequest) (*v1pb.AssistantTestResponse, error) {
	id, e := utils.ExtractIDFromUrn(utils.AssistantUrnPrefix, request.Urn)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to extract id from urn: %v", e)
	}

	assistant, e := s.Gds.Handler.GetLLMAssistant(ctx, id)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary: %v", e)
	}

	res := &v1pb.AssistantTestResponse{
		Tid: shortuuid.New(),
	}

	if request.InteractType == cpb.InteractType_CHAT {
		s.registerAssistantWebsocket(ctx, assistant, request.RagMode, request.Content, request.LlmModel, res.Tid)
	} else if request.InteractType == cpb.InteractType_QUERY {
		return s.ProcessCompleteAssistantTest(ctx, assistant, request.RagMode, request.Content, request.LlmModel)
	}

	return res, nil
}

func (s *APIV1Service) ProcessCompleteAssistantTest(ctx context.Context, srcAssistant *model.LlmAssistant, ragMode cpb.RagMode, content string, modelUrn string) (*v1pb.AssistantTestResponse, error) {
	p := process.NewProcess(s.Gds)
	modelId := srcAssistant.ModelID
	if modelUrn != "" {
		i, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, modelUrn)
		if err == nil {
			modelId = i
		}
	}

	lp, lm, err := p.GetModelProvider(ctx, modelId)
	if err != nil {
		slog.Error("Failed to get model provider", slog.Any("error", err))
		return nil, status.Errorf(codes.Internal, "failed to get model provider: %v", err)
	}

	llmConfig := &llmMaker.ChatClientConfig{
		Url:   lp.AccessURL,
		Model: lm.Model,
		API:   lp.APIName,
	}

	client := llmMaker.NewChatClient(llmConfig)

	messages := []llm.RoleContent{}
	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleSys,
		Content: srcAssistant.SystemPrompt,
	})

	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleUser,
		Content: content,
	})

	req := &llm.CompleteReq{Msgs: messages}
	answer, err := client.Complete(ctx, req)
	if err != nil {
		slog.Error("failed to  assistant complete test", slog.Any("err", err))
		return nil, status.Errorf(codes.Internal, "failed to  assistant complete test: %v", err)
	}

	if answer == nil {
		return nil, status.Errorf(codes.Internal, "failed to  assistant complete test")
	}

	res := &v1pb.AssistantTestResponse{
		Content: answer.GetContent(),
		Tid:     shortuuid.New(),
	}
	return res, nil
}

func (s *APIV1Service) registerAssistantWebsocket(ctx context.Context, srcAssistant *model.LlmAssistant, ragMode cpb.RagMode, content string, modelUrn string, tid string) {
	s.WsManager.InitMessage[tid] = content
	s.WsManager.Register(tid, func(conn *websocket.Conn) {
		slog.Info("WebSocket connection established for assistant test:", slog.Any("UID", tid))
		outMessageChan := make(chan *v1pb.Exchange, 5) // 缓冲大小可以根据需要调整
		inMessageChan := make(chan *v1pb.Exchange, 5)  // 缓冲大小可以根据需要调整

		nctx, cancel := context.WithCancel(context.Background())
		nctx = context.WithValue(nctx, userIdContextKey, srcAssistant.CreatorID)
		userId, err := s.GetCurrentUserId(ctx) // 获取当前用户ID
		if err != nil {
			slog.Error("Failed to get current user ID", slog.Any("error", err))
			return
		}

		assistantService := &AssistantService{
			// done:      make(chan struct{}),
			userId: userId,
			api:    s,
		}

		defer conn.Close()
		defer close(inMessageChan)  // 确保在连接关闭时关闭输入通道
		defer close(outMessageChan) // 确保在处理完成后关闭 channel
		defer cancel()              // 确保在处理完成后取消 context
		defer s.WsManager.UnRegister(tid)

		go assistantService.processAssistantTestRequest(nctx, srcAssistant, ragMode, modelUrn, inMessageChan, outMessageChan)
		go s.handleCommonWebSocketMessages(nctx, cancel, tid, conn, ragMode, inMessageChan)
		s.handleExchangeMessage(nctx, cancel, tid, conn, outMessageChan)
	})
}

func (s *AssistantService) processAssistantTestRequest(ctx context.Context, srcAssistant *model.LlmAssistant, ragMode cpb.RagMode, modelUrn string, inMessageChan <-chan *v1pb.Exchange, outMessageChan chan<- *v1pb.Exchange) {
	p := process.NewProcess(s.api.Gds)
	modelId := srcAssistant.ModelID
	if modelUrn != "" {
		id, e := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, modelUrn)
		if e == nil {
			modelId = id
		}
	}
	lp, lm, err := p.GetModelProvider(ctx, modelId)
	if err != nil {
		slog.Error("Failed to get model provider", slog.Any("error", err))
		return
	}

	llmConfig := &llmMaker.ChatClientConfig{
		Url:   lp.AccessURL,
		Model: lm.Model,
		API:   lp.APIName,
	}
	if lp.AuthInfo != "" {
		key := []byte(viper.GetString("encryption-key"))
		authInfo, err := utils.Decrypt(lp.AuthInfo, key)
		if err != nil {
			slog.Error("Failed to decrypt auth info", slog.Any("error", err))
			return
		}
		llmConfig.AuthInfo = authInfo
	}

	s.client = llmMaker.NewChatClient(llmConfig)
	if s.client == nil {
		slog.Error("Failed to create chat client")
		return
	}

	s.prompt = srcAssistant.SystemPrompt

	cctx, cancel := context.WithCancel(context.Background())
	s.ctx = cctx
	for {
		select {
		case <-ctx.Done():
			slog.Info("assistant service for chat service is done")
			cancel()
			return
		case inMsg, ok := <-inMessageChan:
			if !ok {
				slog.Info("recv error.")
				return
			}

			if inMsg.Phase == v1pb.Phase_End {
				slog.Info("Received end message:", slog.Any("id", inMsg.AskId))
				cancel()
				cctx, cancel = context.WithCancel(context.Background())
				break
			}

			if inMsg.Content == "" {
				continue
			}

			msgcontent := s.BuildChatMessage(v1pb.Phase_Prepare, "正在分析中...", "")
			s.sendOutMessage(cctx, outMessageChan, msgcontent)

			if inMsg.Mode != cpb.RagMode_RAG_MODE_UNSPECIFIED {
				ragMode = inMsg.Mode
			}

			go func() {
				s.doChat(ragMode, inMsg.Content, outMessageChan)
			}()
		}

	}
}
